package com.fliad.ai.service;

import cn.hutool.core.util.IdUtil;
import com.fliad.ai.core.AiFlowDriver;
import com.fliad.ai.dto.taskmodel.*;
import org.noear.liquor.eval.Exprs;
import org.noear.snack.ONode;
import org.noear.solon.annotation.Component;
import org.noear.solon.core.event.AppLoadEndEvent;
import org.noear.solon.core.event.EventListener;
import org.noear.solon.core.util.RunUtil;
import org.noear.solon.expression.snel.SnEL;
import org.noear.solon.flow.Chain;
import org.noear.solon.flow.FlowContext;
import org.noear.solon.flow.NodeDecl;
import org.noear.solon.flow.evaluation.LiquorEvaluation;
import org.noear.solon.flow.stateful.StatefulFlowEngine;
import org.noear.solon.flow.stateful.controller.BlockStateController;
import org.noear.solon.flow.stateful.repository.InMemoryStateRepository;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class TaskService implements EventListener<AppLoadEndEvent> {

    StatefulFlowEngine statefulFlowEngine;
    private Map<String, FlowContext> contexts = new ConcurrentHashMap<>();

    /**
     * 运行任务
     *
     * @param request
     * @return
     */
    public TaskRunOutput taskRun(TaskRunInput request) {
        FlowContext context = new FlowContext(null);
        String taskID = IdUtil.getSnowflakeNextIdStr();
        // 任务报告
        TaskReportOutput report = new TaskReportOutput(request.getInputs());
        context.put("report", report);
        Map<String, Object> variableContext = new HashMap<>();
        context.put("variableContext", variableContext);
        contexts.put(taskID, context);
        Chain chain = new Chain(IdUtil.getSnowflakeNextIdStr(), "");
        ONode.load(request.getSchema()).get("nodes").forEach(node -> {
            if ("start".equals(node.get("type").getString())) {
                NodeDecl nodeDecl = NodeDecl.startOf(node.get("id").getString()).title(node.get("data").get("title").getString());
                nodeDecl.metaPut("outputs", node.get("data").get("outputs"));
                ONode.load(request.getSchema()).get("edges").forEach(edge -> {
                    if (node.get("id").getString().equals(edge.get("sourceNodeID").getString())) {
                        nodeDecl.linkAdd(edge.get("targetNodeID").getString());
                    }
                });
                chain.addNode(nodeDecl);
                Map<String, Object> startMap = new HashMap<>(report.getInputs());
                variableContext.put(node.get("id").getString(), startMap);
            } else if ("end".equals(node.get("type").getString())) {
                NodeDecl nodeDecl = NodeDecl.endOf(node.get("id").getString()).title(node.get("data").get("title").getString());
                nodeDecl.metaPut("inputs", node.get("data").get("inputs"));
                nodeDecl.metaPut("inputsValues", node.get("data").get("inputsValues"));
                ONode.load(request.getSchema()).get("edges").forEach(edge -> {
                    if (node.get("id").getString().equals(edge.get("sourceNodeID").getString())) {
                        nodeDecl.linkAdd(edge.get("targetNodeID").getString());
                    }
                });
                chain.addNode(nodeDecl);
            } else if ("llm".equals(node.get("type").getString())) {
                NodeDecl nodeDecl = NodeDecl.activityOf(node.get("id").getString()).title(node.get("data").get("title").getString())
                        .task("@llm");
                nodeDecl.metaPut("inputs", node.get("data").get("inputs"));
                nodeDecl.metaPut("inputsValues", node.get("data").get("inputsValues"));
                nodeDecl.metaPut("outputs", node.get("data").get("outputs"));
                ONode.load(request.getSchema()).get("edges").forEach(edge -> {
                    if (node.get("id").getString().equals(edge.get("sourceNodeID").getString())) {
                        nodeDecl.linkAdd(edge.get("targetNodeID").getString());
                    }
                });
                chain.addNode(nodeDecl);
                //  chain.addNode(NodeDecl.of(node.get("id").getString()).title(node.get("data").get("title").getString()));
            } else if ("condition".equals(node.get("type").getString())) {
                // 包容网关
                NodeDecl nodeDecl = NodeDecl.inclusiveOf(node.get("id").getString()).title(node.get("data").get("title").getString());
                nodeDecl.metaPut("conditions", node.get("data").get("conditions"));
                node.get("data").get("conditions").forEach(conditionNode -> {
                    ONode.load(request.getSchema()).get("edges").forEach(edge -> {
                        if (conditionNode.get("key").getString().equals(edge.get("sourcePortID").getString())) {
                            String condition = String.format("%s.%s.%s %s %s", "variableContext", conditionNode.get("value").get("left").get("content").get(0).getString(),
                                    conditionNode.get("value").get("left").get("content").get(1).getString(),
                                    "eq".equals(conditionNode.get("value").get("operator").getString()) ? "==" : "!=",
                                    conditionNode.get("value").get("right").get("content"));
                            nodeDecl.linkAdd(edge.get("targetNodeID").getString(), l -> l.when(condition));
                        }
                    });
                });
                chain.addNode(nodeDecl);
            }
        });

        RunUtil.asyncAndTry(() -> {
            statefulFlowEngine.eval(chain, context);
        });
        return new TaskRunOutput(taskID);
    }

    /**
     * 任务报告
     *
     * @param request
     * @return
     */
    public TaskReportOutput taskReport(TaskReportInput request) {
        return contexts.get(request.getTaskID()).get("report");
    }

    public TaskResultOutput taskResult(TaskResultInput request) {
        return null;
    }

    /**
     * 停止任务
     *
     * @param request
     * @return
     */
    public TaskCancelOutput taskCancel(TaskCancelInput request) {
        // 停止当前实例
        if (contexts.get(request.getTaskID()) != null) {
            contexts.get(request.getTaskID()).stop();
            TaskReportOutput report = contexts.get(request.getTaskID()).get("report");
            report.getWorkflowStatus().setTerminated(true);
        }
        return new TaskCancelOutput(true);
    }

    @Override
    public void onEvent(AppLoadEndEvent appLoadEndEvent) throws Throwable {
        statefulFlowEngine = StatefulFlowEngine.newInstance(AiFlowDriver.builder1()
                .stateController(new BlockStateController())
                .stateRepository(new InMemoryStateRepository()) //状态仓库（支持持久化）
                .evaluation(new LiquorEvaluation() {
                    @Override
                    public boolean runTest(FlowContext context, String code) {
                        return (boolean) SnEL.eval(code, context.model());
                    }
                })
                .build());
    }

    public static void main(String[] args) {
        String code = "variableContext.start_0.input == \"你好\"";
        Map<String, Object> context = new HashMap<>();
        context.put("a", 1);
        context.put("b", 2);
        context.put("input", "你好1");

        Exprs.eval("(a + b) * 2", context);
        System.out.println(Exprs.eval("(a + b) * 2", context));
        System.out.println(SnEL.eval("input != '你好'", context));
    }
}
